#ifndef _AVLTREE_H_
#define _AVLTREE_H_

#include "Common.h"
#include "ICompare.h"
#include "TreeNode.h"
#include "Exceptions.h"
#include "Stack.h"
#include <math.h>
#include <cassert>

namespace DataStructuresWet1
{
template <class DataT>
class AVLTree
{
public:
	AVLTree(ICompare<DataT>& comparator);
	AVLTree(ICompare<DataT>& comparator, int size);
	//AVLTree(const AVLTree& other);
	~AVLTree();
	void Insert(DataT* data);
	void Remove(DataT& data);
	DataT* Find(DataT& data) const;
	DataT* Select(int index) const;
	AVLTree<DataT>* Merge(const AVLTree<DataT>& other);
	int Size() const;
	void ToSortedArray(DataT* arr[]) const;
	void FromSortedArray(DataT* arr[]);
	DataT* MinElement() const;
	DataT* MaxElement() const;
	DataT* MedianElement() const;
private:
	typedef _TreeNode<DataT> TreeNode;
	ICompare<DataT> &Comparator;
	TreeNode* Root;

	void ReBalance(TreeNode* node);
	DataT* DoSelect(TreeNode* node, int index) const;
	void BuildFullSubTree(TreeNode* root, int height);
	void RemoveLeavesFromFullTree(int numOfLeaves);
};


template <class DataT>
AVLTree<DataT>::AVLTree(ICompare<DataT> &comparator) : Comparator(comparator), Root(NULL) {}
//
//template <class DataT>
//AVLTree<DataT>::AVLTree(const AVLTree<DataT> &other) : Comparator(other.Comparator)
//{
//	//TODO: Implement
//}

template <class DataT>
AVLTree<DataT>::AVLTree(ICompare<DataT> &comparator, int size) : Comparator(comparator), Root(NULL)
{
	int height = (int) ceil(log2((float) this->Size()));
	int diff = (int)(pow((float) 2, height)) - this->Size();
	TreeNode* root = new TreeNode(NULL);
	BuildFullSubTree(root, height - 1);
	RemoveLeavesFromFullTree(diff);
}

template <class DataT>
AVLTree<DataT>::~AVLTree() {}

template <class DataT>
DataT* AVLTree<DataT>::MedianElement() const
{

	int medianIndex = (int) ceil((float) (this->Size()) / 2) - 1;
	rreturn this->Select(medianIndex);
}

template <class DataT>
DataT* AVLTree<DataT>::MinElement() const
{
	return Root->GetLeftMostChild()->GetDataPtr();
}

template <class DataT>
DataT* AVLTree<DataT>::MaxElement() const
{
	return Root->GetRightMostChild()->GetDataPtr();
}

template <class DataT>
int AVLTree<DataT>::Size() const
{
	if (Root == NULL)
	{
		return 0;
	}
	return Root->GetSize();
}

template <class DataT>
void AVLTree<DataT>::BuildFullSubTree(typename AVLTree<DataT>::TreeNode* root, int height)
{
	if (height = 0)
	{
		return;
	}
	TreeNode* leftChild = new TreeNode(NULL);
	BuildFullSubTree(leftChild, height - 1);
	root->SetLeftChild(leftChild);
	TreeNode* rightChild = new TreeNode(NULL);
	BuildFullSubTree(rightChild, height - 1);
	root->SetRightChild(rightChild);
}

template <class DataT>
void AVLTree<DataT>::RemoveLeavesFromFullTree(int numOfLeaves)
{
	TreeNode* node = this->Root;
	Stack<TreeNode> s (node->GetHeight());
	int size = node->GetSize();
	for (int i = 0; i < numOfLeaves;)
	{
		while (node != NULL)
		{
			s.Push(node);
			node = node->GetRightChild();
		}
	
		node = s.Pop();
		if (node->NodeChildrenType() == _TreeNode<DataT>::LEAF)
		{
			node->MakeOrphan();
			delete node;
			i++;
			node = NULL; //for next iteration
		}
		node = node->GetLeftChild();
	}
}


template <class DataT>
void AVLTree<DataT>::ToSortedArray(DataT* arr[]) const
{
	TreeNode* node = this->Root;
	Stack<TreeNode> s (node->GetHeight());
	int size = node->GetSize();
	for (int i = 0; i < size; i++)
	{
		while (node != NULL)
		{
			s.Push(node);
			node = node->GetLeftChild();
		}
		node = s.Pop();
		arr[i] = node->GetDataPtr();
		node = node->GetRightChild();
	}
	assert(s.Pop() == NULL); //stack is empty
}

template <class DataT>
void AVLTree<DataT>::FromSortedArray(DataT* arr[])
{
	TreeNode* node = this->Root;
	Stack<TreeNode> s (node->GetHeight());
	int size = node->GetSize();
	for (int i = 0; i < size; i++)
	{
		while (node != NULL)
		{
			s.Push(node);
			node = node->GetLeftChild();
		}
		node = s.Pop();
		node->SetData(arr[i]);
		node = node->GetRightChild();
	}
	assert(s.Pop() == NULL); //stack is empty
}

template <class DataT>
DataT* AVLTree<DataT>::Find(DataT& data) const
{
	//TODO: Implement
	return &data;
}

template <class DataT>
void AVLTree<DataT>::Insert(DataT* data)
{
	TreeNode* node = this->Root;
	if (node == NULL)
	{
		this->Root = new TreeNode(data);
		return;
	}
	while(node != NULL)
	{
		if (*node->GetDataPtr() == *data)
		{
			throw FailureException();
		}
		if (Comparator.IsLeq(*data, *(node->GetDataPtr())))
		{
			if (node->GetLeftChild() == NULL)
			{
				node->SetLeftChild(new TreeNode(data));
				break;
			}
		}
		else
		{
			if (node->GetRightChild() == NULL)
			{
				node->SetRightChild(new TreeNode(data));
				break;
			}
		}
	}
	assert(node != NULL);
	while(node != NULL)
	{
		ReBalance(node);
		node = node->GetParent();
	}
}


template <class DataT>
void AVLTree<DataT>::Remove(DataT& data)
{
	TreeNode* node = this->Root;
	TreeNode* ReBalanceFrom;
	if (node == NULL)
	{
		throw FailureException();
	}
	while(node != NULL)
	{
		if (Comparator.IsEqual(data, *(node->GetDataPtr())))
		{
			switch(node->NodeChildrenType())
			{
			case _TreeNode<DataT>::LEAF:
				ReBalanceFrom = node->GetParent();
				node->MakeOrphan();
				break;
			case _TreeNode<DataT>::HAS_LEFT_CHILD:
				ReBalanceFrom = node->GetParent();
				node->GetLeftChild()->ReplaceParent();
				break;
			case _TreeNode<DataT>::HAS_RIGHT_CHILD:
				ReBalanceFrom = node->GetParent();
				node->GetRightChild()->ReplaceParent();
				break;
			case _TreeNode<DataT>::HAS_TWO_CHILDREN:
				node->SwitchPosition(*(node->GetRightChild()->GetLeftMostChild()));
				ReBalanceFrom = node->GetParent();
				if (node->NodeChildrenType() == _TreeNode<DataT>::HAS_RIGHT_CHILD)
				{
					node->GetRightChild()->ReplaceParent();
				}
				else
				{
					node->MakeOrphan();
				}
				break;
			}
			delete node;
		}
	node = ReBalanceFrom;
	while(node != NULL)
	{
		ReBalance(node);
		node = node->GetParent();
	}
	}
}

template <class DataT>
AVLTree<DataT>* AVLTree<DataT>::Merge(const AVLTree<DataT>& other)
{
	int size1 = this->Size();
	int size2 = other.Size();
	if (size1 == 0)
	{
		return new AVLTree<DataT>(other);
	}
	if (size2 == 0)
	{
		return new AVLTree<DataT>(*this);
	}

	DataT** arr1 = new DataT*[size1];
	DataT** arr2 = new DataT*[size2];
	this->ToSortedArray(arr1);
	other.ToSortedArray(arr2);
	int newSize = size1 + size2;
	DataT** newArr = new DataT*[newSize];
	MergeArrays(arr1, size1, arr2, size2, newArr);
	AVLTree<DataT>* newTree = new AVLTree(this->Comparator, newSize);
	newTree->FromSortedArray(newArr);
	delete [] arr1;
	delete [] arr2;
	delete [] newArr;
	return newTree;
}

template <class DataT>
void AVLTree<DataT>::ReBalance(TreeNode* node)
{
	node->UpdateSize();
	node->UpdateHeight();
	int nodeBF = node->GetBalanceFactor();
	int childBF;
	assert(abs(nodeBF) <= 2);
	if (nodeBF == 2)
	{
		childBF = node->GetLeftChild()->GetBalanceFactor();
		assert(abs(childBF) <= 2);
		if (childBF >= 0)
		{
			node->RotateLL();
		}
		else //childBF == -1
		{
			node->RotateLR();
		}
	}
	if (nodeBF == -2)
	{
		childBF = node->GetRightChild()->GetBalanceFactor();
		assert(abs(childBF) <= 2);
		if (childBF <= 0)
		{
			node->RotateRR();
		}
		else //childBF == 1
		{
			node->RotateRL();
		}
	}
}

template <class DataT>
DataT* AVLTree<DataT>::Select(int index) const
{
	TreeNode* root = this->Root;
	return DoSelect(root, index);
}

template <class DataT>
DataT* AVLTree<DataT>::DoSelect(typename AVLTree<DataT>::TreeNode* node, int index) const
{
	if (node == NULL || node->GetWeight() == index)
	{
		return node;
	}

	if (node->GetWeight() > index)
	{
		DoSelect(node->GetLeftChild(), index);
	}
	else //node->GetWeight() < index
	{
		DoSelect(node->GeRightChild(), index - node->GetWeight());
	}
}

} //DataStructuresWet1


#endif //_AVLTREE_H_